import os
from torch.utils.data import DataLoader
from torchvision import datasets


class MyDataLoader(DataLoader):
    def __init__(self, path, transform_=None, batch_size_=1, shuffle_=False, num_workers_=4,
                 pin_memory_: bool = True):
        """
        自定义DataLoader
        :param path: 训练集、测试集、验证集路径
        :param transform_: 数据变换
        :param batch_size_: 批处理数量
        :param shuffle_: 是否打乱dataset
        :param num_workers_: 加载数据集的子线程数
        :param pin_memory_: 是否固定内存
        """
        assert os.path.exists(path), "{} path does not exist!".format(path)

        # init
        self.__path = path  # 数据集路径
        self.__dataset = datasets.ImageFolder(path, transform=transform_)  # 创建dataset对象
        self.__size = len(self.__dataset)  # 数据集尺寸
        self.__class_dict = self.__dataset.class_to_idx  # 类别名字典{['key': val],...}

        # 继承父类，并覆盖父类的init方法，在init中添加父类的init方法
        super(MyDataLoader, self).__init__(self.__dataset, batch_size=batch_size_, shuffle=shuffle_,
                                           num_workers=num_workers_, pin_memory=pin_memory_)

    @property
    def size(self):
        return self.__size

    @property
    def class_dict(self):
        return self.__class_dict
